import csv
import sys
import os
from io import BytesIO

import time
import unittest

from rdbtools import RdbParser
from rdbtools import MemoryCallback


from rdbtools.memprofiler import MemoryRecord, PrintAllKeys

def format_bytes(size_bytes):
    """
    将字节数自动转换为合适的单位，如 B, KB, MB, GB, TB，保留两位小数。

    """
    # 定义单位
    units = ['B', 'KB', 'MB', 'GB', 'TB', 'PB']
    if size_bytes == 0:
        return "0.00 B"
    # 计算单位索引
    i = 0
    while size_bytes >= 1024 and i < len(units) - 1:
        size_bytes /= 1024
        i += 1
    # 保留两位小数
    return f"{size_bytes:.2f} {units[i]}"

class Stats(object):
    def __init__(self,dump_name,csv_filename="output.csv",csv_mode = "w"):
        self.sums = {} #统计各类型的总大小
        self.records = {} 
        self.permanent_size = 0
        self.all_keys_size = 0
        self.all_keys = 0 #key总数
        self.permanent_count = 0 #永久key总数
        self.shardname = dump_name
        #操作csv写入
        self.csv_filename = csv_filename
        self.csv_file = open(self.csv_filename, mode=csv_mode, newline='', encoding='utf-8')
        self.writer = csv.writer(self.csv_file)
        if csv_mode == "w":
          self.writer.writerow(["RDB",'DB', 'Key', 'Type', 'Size (Bytes)', 'Expiry'])


    def next_record(self, record):
        self.all_keys +=1 
        self.all_keys_size += record.bytes
        if record.type not in self.sums:
            self.sums.setdefault(record.type,0)
        self.sums[record.type] += record.bytes
        is_permanent = record.expiry in (None, -1, 0)
        if record.key is not None:
          if is_permanent:
              self.permanent_count += 1
              self.permanent_size += record.bytes
              self.writer.writerow((self.shardname,record.database, record.key, record.type, record.bytes, record.expiry))

    def close(self):
        # 关闭 CSV 文件
        self.csv_file.close()
        

def get_stats(files):
    for index,file_name in enumerate(files.split(",")):
      if index != 0:
        stats = Stats(os.path.splitext(os.path.basename(file_name))[0],csv_mode="a")
      else:
        stats = Stats(os.path.splitext(os.path.basename(file_name))[0])
      callback = MemoryCallback(stats, 64)
      parser = RdbParser(callback)
      parser.parse(file_name)
    return stats


def get_sums(file_name):
    stats = Stats()
    callback = MemoryCallback(stats, 64)
    parser = RdbParser(callback)
    parser.parse(os.path.join(os.path.dirname(__file__), 'dumps', file_name))
    return stats.sums

def get_csv(dump_file_name):
    buff = BytesIO()
    callback = MemoryCallback(PrintAllKeys(buff, None, None), 64)
    parser = RdbParser(callback)
    parser.parse(dump_file_name)
    csv = buff.getvalue().decode()
    return csv

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("🔧 示例: python memprofiler.py ./dump.rdb")
        sys.exit(1)

    rdb_file_path = sys.argv[1]

    print(f"🔍 正在解析 RDB 文件: {rdb_file_path}")
    start_time = time.time()

    stat = get_stats(rdb_file_path)
    print(f"\n✅ 分析完成.")
    elapsed_time = time.time() - start_time
    print(f"🕒 分析耗时：{elapsed_time:.2f} 秒") 
    print(f"🔢 总 Key 数量: {stat.all_keys} 内存大小: {format_bytes(stat.all_keys_size)}")
    print(f"🔒 永久 Key（无过期时间）数量: {stat.permanent_count} 内存大小: {format_bytes(stat.permanent_size)}")
    print(f"📄 CSV 文件已保存为: {stat.csv_filename}")